package embedding;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Collectors;

import com.fasterxml.jackson.databind.ObjectMapper;

/**
 * OpenAI 嵌入向量生成器实现
 */
public class OpenAIEmbedding extends Embedding {
    private String apiKey;
    private String model;
    private String baseURL;
    private int dimension;
    private HttpClient client;
    private ObjectMapper objectMapper;

    /**
     * OpenAI 嵌入配置类
     */
    public static class OpenAIEmbeddingConfig {
        private String model;
        private String apiKey;
        private String baseURL;

        public OpenAIEmbeddingConfig(String apiKey, String model) {
            this.apiKey = apiKey;
            this.model = model;
        }

        public OpenAIEmbeddingConfig(String apiKey, String model, String baseURL) {
            this.apiKey = apiKey;
            this.model = model;
            this.baseURL = baseURL;
        }

        public String getModel() {
            return model;
        }

        public void setModel(String model) {
            this.model = model;
        }

        public String getApiKey() {
            return apiKey;
        }

        public void setApiKey(String apiKey) {
            this.apiKey = apiKey;
        }

        public String getBaseURL() {
            return baseURL;
        }

        public void setBaseURL(String baseURL) {
            this.baseURL = baseURL;
        }
    }

    /**
     * 创建一个新的 OpenAI 嵌入向量生成器
     * 
     * @param config OpenAI 嵌入配置
     */
    public OpenAIEmbedding(OpenAIEmbeddingConfig config) {
        this.apiKey = config.getApiKey();
        this.model = config.getModel() != null ? config.getModel() : "text-embedding-3-small";
        this.baseURL = config.getBaseURL() != null ? config.getBaseURL() : "https://api.openai.com";
        this.maxTokens = 8192; // OpenAI 嵌入模型的最大标记数
        this.client = HttpClient.newBuilder()
                .connectTimeout(Duration.ofSeconds(30))
                .build();
        this.objectMapper = new ObjectMapper();
        
        // 根据模型设置维度
        updateDimensionForModel(this.model);
    }

    /**
     * 根据模型更新维度
     * 
     * @param model 模型名称
     */
    private void updateDimensionForModel(String model) {
        if ("text-embedding-3-small".equals(model)) {
            this.dimension = 1536;
        } else if ("text-embedding-3-large".equals(model)) {
            this.dimension = 3072;
        } else if ("text-embedding-ada-002".equals(model)) {
            this.dimension = 1536;
        } else {
            this.dimension = 1536; // 默认维度
        }
    }

    @Override
    public CompletableFuture<EmbeddingVector> embed(String text) {
        String processedText = preprocessText(text);
        
        return CompletableFuture.supplyAsync(() -> {
            try {
                Map<String, Object> requestBody = new HashMap<>();
                requestBody.put("model", this.model);
                requestBody.put("input", processedText);
                requestBody.put("encoding_format", "float");
                
                String requestBodyJson = objectMapper.writeValueAsString(requestBody);
                
                HttpRequest request = HttpRequest.newBuilder()
                        .uri(URI.create(this.baseURL + "/v1/embeddings"))
                        .header("Content-Type", "application/json")
                        .header("Authorization", "Bearer " + this.apiKey)
                        .POST(HttpRequest.BodyPublishers.ofString(requestBodyJson))
                        .build();
                
                HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
                
                if (response.statusCode() != 200) {
                    throw new RuntimeException("OpenAI API error: " + response.body());
                }
                
                Map<String, Object> responseMap = objectMapper.readValue(response.body(), Map.class);
                List<Map<String, Object>> data = (List<Map<String, Object>>) responseMap.get("data");
                List<Double> embedding = (List<Double>) data.get(0).get("embedding");
                
                double[] vector = embedding.stream().mapToDouble(Double::doubleValue).toArray();
                return new EmbeddingVector(vector, this.dimension);
            } catch (Exception e) {
                throw new RuntimeException("Failed to generate embedding: " + e.getMessage(), e);
            }
        });
    }

    @Override
    public CompletableFuture<List<EmbeddingVector>> embedBatch(List<String> texts) {
        List<String> processedTexts = preprocessTexts(texts);
        
        return CompletableFuture.supplyAsync(() -> {
            try {
                Map<String, Object> requestBody = new HashMap<>();
                requestBody.put("model", this.model);
                requestBody.put("input", processedTexts);
                requestBody.put("encoding_format", "float");
                
                String requestBodyJson = objectMapper.writeValueAsString(requestBody);
                
                HttpRequest request = HttpRequest.newBuilder()
                        .uri(URI.create(this.baseURL + "/v1/embeddings"))
                        .header("Content-Type", "application/json")
                        .header("Authorization", "Bearer " + this.apiKey)
                        .POST(HttpRequest.BodyPublishers.ofString(requestBodyJson))
                        .build();
                
                HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
                
                if (response.statusCode() != 200) {
                    throw new RuntimeException("OpenAI API error: " + response.body());
                }
                
                Map<String, Object> responseMap = objectMapper.readValue(response.body(), Map.class);
                List<Map<String, Object>> data = (List<Map<String, Object>>) responseMap.get("data");
                
                List<EmbeddingVector> embeddings = new ArrayList<>();
                for (Map<String, Object> item : data) {
                    List<Double> embedding = (List<Double>) item.get("embedding");
                    double[] vector = embedding.stream().mapToDouble(Double::doubleValue).toArray();
                    embeddings.add(new EmbeddingVector(vector, this.dimension));
                }
                
                return embeddings;
            } catch (Exception e) {
                throw new RuntimeException("Failed to generate batch embeddings: " + e.getMessage(), e);
            }
        });
    }

    @Override
    public int getDimension() {
        return this.dimension;
    }

    @Override
    public String getProvider() {
        return "OpenAI";
    }

    /**
     * 设置模型类型
     * 
     * @param model 模型名称
     */
    public void setModel(String model) {
        this.model = model;
        updateDimensionForModel(model);
    }

    /**
     * 获取支持的模型列表
     * 
     * @return 支持的模型列表
     */
    public static Map<String, Map<String, Object>> getSupportedModels() {
        Map<String, Map<String, Object>> models = new HashMap<>();
        
        Map<String, Object> small = new HashMap<>();
        small.put("dimension", 1536);
        small.put("description", "High performance and cost-effective embedding model (recommended)");
        models.put("text-embedding-3-small", small);
        
        Map<String, Object> large = new HashMap<>();
        large.put("dimension", 3072);
        large.put("description", "Highest performance embedding model with larger dimensions");
        models.put("text-embedding-3-large", large);
        
        Map<String, Object> ada = new HashMap<>();
        ada.put("dimension", 1536);
        ada.put("description", "Legacy model (use text-embedding-3-small instead)");
        models.put("text-embedding-ada-002", ada);
        
        return models;
    }
}